'''
{
    data:[
        {
            img_path:
            depth_path:
            localimg_emb_path:
        }
    ]
}
'''
import os
from os.path import exists as ope
from os.path import join as osj
from tqdm import tqdm
import json,pdb
save_path = 'local_data_for_depth.json'

types = ['collar','sleeve','pockets']
will_save = {'data':[]}
# choose_type = types[2]
for choose_type in types:
    img_dir = f'/mnt/nas/shengjie/datasets/cloth_{choose_type}_balanced'
    depth_dir = f'/mnt/nas/shengjie/datasets/cloth_{choose_type}_balanced_depth'
    localimg_emb_dir = f'/mnt/nas/shengjie/datasets/cloth_{choose_type}_localimg_emb'
    assert ope(img_dir),img_dir
    assert ope(depth_dir),depth_dir
    assert ope(localimg_emb_dir),localimg_emb_dir

    count = 0
    for entry in os.scandir(img_dir):
    # for filename in  tqdm(os.listdir(img_dir)):
        filename = entry.name
        if not filename.endswith('.jpg'):continue
        count += 1
        # if count % 5==0:
        #     print(count,end=' ',flush=True)
        # if count % 1000 == 0:
        #     print()
        imgname = filename
        depthname = filename.replace('.jpg','.png')
        localimg_emb_name = filename.replace('.jpg','.pth')
        img_path = osj(img_dir,imgname)
        depth_path = osj(depth_dir,depthname)
        localimg_emb_path = osj(localimg_emb_dir,localimg_emb_name)
        assert ope(img_path),img_path
        assert ope(depth_path),depth_path
        assert ope(localimg_emb_path),localimg_emb_path

        will_save['data'].append({
            'img_path':img_path,
            'depth_path':depth_path,
            'localimg_emb_path':localimg_emb_path,
        })
        # pdb.set_trace()
    print()
    print(f'{choose_type} process num : ',count)
    pdb.set_trace()
with open(save_path,'w',encoding='utf-8') as f:
    json.dump(will_save,f,ensure_ascii=False,indent=4)